#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
文件搜索演示脚本
功能：遍历指定目录，将文件名向量化存储到Milvus，支持语义搜索
"""

from tools.FileSearchTools import FileSearchTool
import os
import time

# 尝试导入配置文件
try:
    from config import (
        MILVUS_HOST, MILVUS_PORT, MILVUS_USERNAME, 
        MILVUS_PASSWORD, MILVUS_TOKEN, TARGET_DIRECTORY,
        MAX_FILES, COLLECTION_NAME
    )
    print("已加载配置文件 config.py")
except ImportError:
    print("未找到 config.py 文件，使用默认配置")
    print("提示: 复制 config.py.example 为 config.py 并配置您的信息")
    
    # 默认配置
    MILVUS_HOST = "c-f88be548ab3aa84d.milvus.aliyuncs.com"  # 不要包含http://前缀
    MILVUS_PORT = 19530
    MILVUS_USERNAME = None
    MILVUS_PASSWORD = None
    MILVUS_TOKEN = None
    TARGET_DIRECTORY = "d:/zhang/Documents/workspace"
    MAX_FILES = 5000
    COLLECTION_NAME = "workspace_files"


def main():
    """主函数"""
    print("=" * 60)
    print("文件搜索系统演示")
    print("=" * 60)
    
    # 使用配置参数
    target_directory = TARGET_DIRECTORY
    max_files = MAX_FILES
    
    print(f"当前配置:")
    print(f"  Milvus服务器: {MILVUS_HOST}:{MILVUS_PORT}")
    print(f"  目标目录: {target_directory}")
    print(f"  文件数量限制: {max_files}")
    
    # 检查目录是否存在
    if not os.path.exists(target_directory):
        print(f"错误: 目录不存在 - {target_directory}")
        print("请检查路径是否正确")
        return
    
    # 检查Milvus认证配置
    username = MILVUS_USERNAME
    password = MILVUS_PASSWORD
    token = MILVUS_TOKEN
    
    if not username or not password:
        if not token:
            print("\n警告: 未设置Milvus认证信息")
            print("请在 config.py 中设置用户名密码或访问令牌")
            
            # 提示用户输入认证信息
            username = input("请输入Milvus用户名 (留空跳过): ").strip()
            if username:
                password = input("请输入Milvus密码: ").strip()
            else:
                token = input("请输入Milvus访问令牌 (留空跳过): ").strip()
            
            if not (username and password) and not token:
                print("未提供认证信息，无法连接到Milvus")
                return
    
    try:
        # 初始化文件搜索工具
        print("\n1. 初始化文件搜索工具...")
        search_tool = FileSearchTool(
            collection_name=COLLECTION_NAME,
            milvus_host=MILVUS_HOST,
            milvus_port=MILVUS_PORT,
            username=username,
            password=password,
            token=token
        )
        
        # 获取当前集合统计
        stats = search_tool.get_collection_stats()
        print(f"当前集合状态: {stats}")
        
        # 询问是否重新索引
        if stats['total_entities'] > 0:
            choice = input(f"\n集合中已有 {stats['total_entities']} 条记录，是否重新索引？(y/n): ").lower()
            if choice == 'y':
                print("正在清空集合...")
                search_tool.clear_collection()
        
        # 如果集合为空，进行索引
        if search_tool.get_collection_stats()['total_entities'] == 0:
            print(f"\n2. 扫描目录: {target_directory}")
            print(f"   最大文件数: {max_files}")
            
            # 扫描目录
            start_time = time.time()
            file_list = search_tool.scan_directory(target_directory, max_files)
            scan_time = time.time() - start_time
            
            if not file_list:
                print("未找到任何文件，请检查目录路径和权限")
                return
            
            print(f"扫描完成，用时 {scan_time:.2f}秒")
            
            # 索引文件
            print(f"\n3. 开始向量化索引...")
            start_time = time.time()
            search_tool.index_files(file_list, batch_size=50)
            index_time = time.time() - start_time
            
            print(f"索引完成，用时 {index_time:.2f}秒")
        else:
            print(f"\n使用现有索引，共 {stats['total_entities']} 条记录")
        
        # 交互式搜索
        print(f"\n4. 开始交互式搜索")
        print("=" * 40)
        print("输入搜索关键词，输入 'quit' 退出")
        print("搜索示例:")
        print("  - python")
        print("  - 配置文件")  
        print("  - 文档")
        print("  - 图片")
        print("=" * 40)
        
        while True:
            query = input("\n请输入搜索关键词: ").strip()
            
            if query.lower() == 'quit':
                print("退出搜索")
                break
            
            if not query:
                print("请输入有效的搜索关键词")
                continue
            
            # 执行搜索
            start_time = time.time()
            results = search_tool.search_files(query, top_k=10)
            search_time = time.time() - start_time
            
            # 显示结果
            print(f"\n搜索结果 (用时 {search_time:.3f}秒):")
            print("-" * 80)
            
            if results:
                for i, result in enumerate(results, 1):
                    file_size_mb = result['file_size'] / 1024 / 1024 if result['file_size'] > 0 else 0
                    print(f"{i:2d}. 文件名: {result['file_name']}")
                    print(f"     路径: {result['file_path']}")
                    print(f"     类型: {result['file_type']:<15} 大小: {file_size_mb:.2f}MB")
                    print(f"     相似度: {result['similarity_score']:.4f}")
                    print(f"     修改时间: {result['modified_time']}")
                    print()
            else:
                print("未找到相关文件")
            
            print("-" * 80)
    
    except KeyboardInterrupt:
        print("\n\n用户中断程序")
    except Exception as e:
        print(f"\n程序执行出错: {str(e)}")
        print("请检查Milvus服务是否正常运行")
    
    print("\n程序结束")


def test_search_examples():
    """测试搜索示例"""
    try:
        search_tool = FileSearchTool(collection_name="workspace_files")
        
        # 测试搜索
        test_queries = [
            "python脚本",
            "配置文件", 
            "文档",
            "图片",
            "工具"
        ]
        
        for query in test_queries:
            print(f"\n搜索: '{query}'")
            results = search_tool.search_files(query, top_k=5)
            
            for i, result in enumerate(results, 1):
                print(f"{i}. {result['file_name']} (相似度: {result['similarity_score']:.4f})")
    
    except Exception as e:
        print(f"测试搜索失败: {e}")


if __name__ == "__main__":
    main()
    
    # 如果需要运行测试，取消下面的注释
    # test_search_examples()